#!/usr/bin/env python3
import rclpy  # ROS 2 Python 客户端库
from rclpy.node import Node  # ROS 2 节点基类
from sensor_msgs.msg import PointCloud2, PointField  # ROS 2 点云消息类型
import numpy as np  # 用于数值计算
import struct  # 用于二进制数据的打包和解包
import open3d as o3d  # 用于3D点云处理和可视化
import time  # 用于计时
import multiprocessing  # 用于并行处理以提高性能

# -----------------------------------------------------------------------------
# 函数：手动将PointCloud2消息解析为NumPy数组
# -----------------------------------------------------------------------------
def pointcloud2_to_numpy(msg: PointCloud2) -> np.ndarray:
    """
    将 ROS2 的 PointCloud2 消息转换为 NumPy 数组，包含 (x, y, z, intensity)。
    这个函数实现了手动解析，避免了对 ros2-numpy 的依赖，提高了代码的通用性。
    """
    # 获取点云消息的元数据
    num_points = msg.width * msg.height  # 总点数
    point_step = msg.point_step  # 每个点占用的字节数
    data = msg.data  # 包含所有点数据的字节数组
    
    # 创建字典，用于存储每个字段的字节偏移量和数据类型
    offsets = {}
    dtypes = {
        PointField.INT8: 'b', PointField.UINT8: 'B',
        PointField.INT16: 'h', PointField.UINT16: 'H',
        PointField.INT32: 'i', PointField.UINT32: 'I',
        PointField.FLOAT32: 'f', PointField.FLOAT64: 'd',
    }
    
    # 遍历 PointCloud2 消息的字段，记录每个字段的偏移量和数据类型
    for field in msg.fields:
        if field.name in ['x', 'y', 'z', 'intensity', 'reflectivity']:
            offsets[field.name] = field.offset  # 存储字段的字节偏移量
            dtypes[field.name] = dtypes[field.datatype]  # 存储字段的数据类型对应的格式符
    
    # 确保必须的 x, y, z 字段存在
    if not all(k in offsets for k in ['x', 'y', 'z']):
        raise ValueError("PointCloud2消息缺少x, y, 或 z字段")
    
    # 确定是否存在强度或反射率字段。如果没有，则使用默认值0
    has_intensity = 'intensity' in offsets or 'reflectivity' in offsets
    intensity_key = 'intensity' if 'intensity' in offsets else 'reflectivity'  # 确定使用哪个字段来获取强度
    
    # 创建一个列表，用于存储每个点的 (x, y, z, intensity) 值
    points = []
    for i in range(num_points):
        # 计算当前点在字节数组中的起始位置
        base_idx = i * point_step
        # 使用 struct.unpack_from() 解包 x, y, z, intensity 的值
        x = struct.unpack_from(dtypes['x'], data, offset=base_idx + offsets['x'])[0]
        y = struct.unpack_from(dtypes['y'], data, offset=base_idx + offsets['y'])[0]
        z = struct.unpack_from(dtypes['z'], data, offset=base_idx + offsets['z'])[0]
        
        # 获取 intensity (或 reflectivity)，如果没有则默认为0
        intensity = 0.0
        if has_intensity:
            intensity = struct.unpack_from(dtypes[intensity_key], data, offset=base_idx + offsets[intensity_key])[0]
        
        # 将 (x, y, z, intensity) 作为一个列表添加到 points 列表中
        points.append([x, y, z, intensity])
    
    # 将点列表转换为 NumPy 数组
    return np.array(points, dtype=np.float32)

# -----------------------------------------------------------------------------
# 函数：创建PointCloud2消息（用于发布）
# -----------------------------------------------------------------------------
# 定义输出PointCloud2消息的字段信息
FIELDS_OUT = [
    PointField(name='x', offset=0, datatype=PointField.FLOAT32, count=1),
    PointField(name='y', offset=4, datatype=PointField.FLOAT32, count=1),
    PointField(name='z', offset=8, datatype=PointField.FLOAT32, count=1),
    PointField(name='intensity', offset=12, datatype=PointField.FLOAT32, count=1),
]
POINT_STEP_OUT = 16  # 输出消息中每个点的字节数

def create_pointcloud2_msg(header, points_array):
    """
    根据给定的 NumPy 数组创建 PointCloud2 消息。
    此函数将NumPy数组中的点云数据转换为ROS可发布的消息格式。
    """
    pcl_msg = PointCloud2()  # 创建一个PointCloud2消息对象
    pcl_msg.header = header  # 设置消息头（时间戳和坐标系）
    pcl_msg.height = 1  # 设置点云高度为1 (无序点云)
    pcl_msg.width = len(points_array)  # 设置点云宽度，即点的数量
    pcl_msg.fields = FIELDS_OUT  # 设置消息的字段定义 (x, y, z, intensity)
    pcl_msg.is_bigendian = False  # 设置字节序 (little endian)
    pcl_msg.point_step = POINT_STEP_OUT  # 设置每个点的字节数
    pcl_msg.row_step = POINT_STEP_OUT * pcl_msg.width  # 设置每行的字节数
    pcl_msg.is_dense = False  # 标记是否包含无效点（这里设为False,因为点云可能包含无效点）
    
    # 将点云数据打包成二进制格式的字节数组
    buffer = bytearray(pcl_msg.row_step) # 创建一个字节数组缓冲区
    for i, point in enumerate(points_array): # 遍历点云中的每一个点
        # 使用 struct.pack_into() 将 x, y, z, intensity 打包到字节数组中
        # 'ffff'：表示四个32位浮点数
        struct.pack_into('ffff', buffer, i * POINT_STEP_OUT, point[0], point[1], point[2], point[3])
    
    pcl_msg.data = buffer  # 设置点云数据
    return pcl_msg  # 返回创建好的PointCloud2消息

# -----------------------------------------------------------------------------
# 函数：并行处理，执行点云的过滤
# -----------------------------------------------------------------------------
def process_chunk(chunk, distance_threshold, exclude_front, area_half_width, area_half_height):
    """
    过滤点云数据。
    """
    # 1. 距离过滤：根据与原点的距离进行过滤
    squared_distances = np.sum(chunk[:, :3] **2, axis=1)  # 计算每个点到原点的距离的平方
    distance_mask = squared_distances <= (distance_threshold** 2)  # 创建距离掩码，只保留在阈值内的点
    
    # 2. 前方区域排除过滤 (如果启用)
    if exclude_front:
        in_front_mask = (chunk[:, 0] > 0)  # 判断点是否在前方
        in_width_mask = (np.abs(chunk[:, 1]) <= area_half_width)  # 判断点是否在宽度范围内
        in_height_mask = (np.abs(chunk[:, 2]) <= area_half_height)  # 判断点是否在高度范围内
        front_area_mask = np.logical_and(np.logical_and(in_front_mask, in_width_mask), in_height_mask)  # 组合三个条件得到前方区域掩码
        filter_mask = np.logical_and(distance_mask, ~front_area_mask)  # 最终过滤掩码：在距离内，但不在前方区域内
    else:
        # 如果不排除前方区域，则只使用距离过滤
        filter_mask = distance_mask
    
    # 使用掩码过滤点云数据，返回过滤后的点云块
    return chunk[filter_mask]

# -----------------------------------------------------------------------------
# 类：LiDAR 点云处理节点
# -----------------------------------------------------------------------------
class LidarProcessor(Node):
    def __init__(self):
        """初始化节点"""
        super().__init__('lidar_processor')  # 初始化ROS2节点，节点名称为 "lidar_processor"
        
        # --- ROS相关设置 ---
        # 声明参数: 输入点云的话题名称
        self.declare_parameter('input_cloud', '/livox/lidar') # 默认点云话题
        input_topic = self.get_parameter('input_cloud').get_parameter_value().string_value  # 获取参数值
        
        # 创建订阅者，订阅点云话题
        self.sub = self.create_subscription(PointCloud2, input_topic, self.callback, 10)
        
        # 创建发布者，发布处理后的点云
        self.pcl_pub = self.create_publisher(PointCloud2, '/livox/pointcloud2_processed', 10)
        
        # 在终端输出一条消息，表示节点已启动，并且显示订阅和发布的话题
        self.get_logger().info(f"订阅: {input_topic}，发布可视化点云到: /livox/pointcloud2_processed")
        
        # --- Open3D 可视化设置 ---
        self.pcd = o3d.geometry.PointCloud()  # 创建一个Open3D点云对象，用于存储点云数据
        self.vis = o3d.visualization.Visualizer()  # 创建Open3D可视化窗口
        self.vis.create_window()  # 创建可视化窗口
        opt = self.vis.get_render_option()  # 获取渲染选项
        opt.point_size = 2.0  # 设置点云中点的大小
        opt.background_color = np.array([0, 0, 0])  # 设置背景颜色为黑色
        self.is_first_frame = True  # 用于标记是否是第一帧点云，如果是第一帧，需要添加几何体
        
        # --- 可调参数 ---
        self.frame_count = 0  # 帧计数器，用于统计处理的帧数
        self.save_path = "/path/to/save/pointclouds/"  # 设置点云保存路径
        self.save_enabled = False  # 标志位，指示是否启用点云保存功能
        
        # 全局过滤参数
        self.distance_threshold = 50.0  # 距离阈值 (米)，超过此距离的点将被滤除
        self.exclude_front = True  # 是否排除前方区域
        self.area_width = 0  # 前方区域的宽度 (米)
        self.area_height = 0  # 前方区域的高度 (米)
        self.area_half_width = self.area_width / 2  # 前方区域半宽度
        self.area_half_height = self.area_height / 2  # 前方区域半高度
        
        # 【新增】聚类感兴趣区域(ROI)参数
        # 定义一个前向长方体区域进行聚类
        # X: 深度方向 (0到50米)
        # Y: 宽度方向 (-1到1米，总宽2米)
        # Z: 高度方向 (-1到1米，总高2米)
        self.clustering_roi = {
            'x': [0.0, 50.0],  # X轴范围 (深度/前向)
            'y': [-1.0, 1.0],  # Y轴范围 (宽度)
            'z': [-1.0, 1.0]   # Z轴范围 (高度)
        }
        
        # 聚类参数
        self.cluster_eps = 0.2  # DBSCAN 邻域半径 (米)
        self.cluster_min_points = 30  # DBSCAN 最小点数
        self.enable_clustering = True  # 是否启用聚类
        self.voxel_size = 0.05  # 体素大小 (米)，用于降采样
        
        # 性能参数
        self.num_processes = max(1, multiprocessing.cpu_count() - 1)  # 并行处理的进程数。 使用CPU核心数-1，保证系统流畅运行

    def callback(self, msg: PointCloud2):
        """
        点云消息回调函数。
        当接收到点云消息时，此函数将被调用。
        """
        start_time = time.time()  # 记录处理开始时间
        
        # 步骤1: 将 ROS 消息转换为 NumPy 数组
        try:
            points_array = pointcloud2_to_numpy(msg)  # 使用自定义函数解析点云消息
        except (ValueError, KeyError) as e:
            self.get_logger().error(f"解析PointCloud2消息失败: {e}")
            return  # 如果解析失败，则直接返回，不再继续处理
        
        if points_array.shape[0] == 0:
            # 如果接收到的点云为空，则记录警告并返回
            self.get_logger().warn("接收到空点云，跳过处理。")
            return
        num_original_points = len(points_array)  # 记录原始点云的点数

        # 步骤2: 全局过滤 (使用并行计算加速)
        chunks = np.array_split(points_array, self.num_processes)  # 将点云分成多个块，用于并行处理
        with multiprocessing.Pool(processes=self.num_processes) as pool:
            # 使用 multiprocessing.Pool 来并行处理点云块
            #  每个点云块被传递给 process_chunk 函数进行处理
            results = pool.starmap(process_chunk, [(chunk, self.distance_threshold, self.exclude_front, self.area_half_width, self.area_half_height) for chunk in chunks])
        # 将所有处理好的点云块合并成一个点云
        filtered_points = np.concatenate(results) if results else np.array([])
        
        if filtered_points.shape[0] == 0:
            # 如果过滤后没有点，清空Open3D点云，更新窗口，并返回
            self.pcd.clear()  # 清空点云
            self.vis.update_geometry(self.pcd)  # 更新可视化
            self.vis.poll_events(); self.vis.update_renderer()
            return  # 跳过后续处理
        
        # 步骤2.5: 分割点云为ROI和背景
        x_mask = (filtered_points[:, 0] >= self.clustering_roi['x'][0]) & (filtered_points[:, 0] <= self.clustering_roi['x'][1]) # ROI内的x轴范围
        y_mask = (filtered_points[:, 1] >= self.clustering_roi['y'][0]) & (filtered_points[:, 1] <= self.clustering_roi['y'][1]) # ROI内的y轴范围
        z_mask = (filtered_points[:, 2] >= self.clustering_roi['z'][0]) & (filtered_points[:, 2] <= self.clustering_roi['z'][1]) # ROI内的z轴范围
        roi_mask = x_mask & y_mask & z_mask  # 创建一个布尔掩码，用于标记在ROI内的点
        
        # 根据掩码将点云分割为ROI内的点和背景点
        points_for_clustering = filtered_points[roi_mask]
        points_background = filtered_points[~roi_mask] # 背景点（不在ROI内的点）
        
        # 步骤3: 对ROI内的点进行聚类和着色
        final_points = filtered_points  # 最终的点云，初始化为已过滤的点云
        final_colors = np.ones((len(filtered_points), 3)) * 0.5  # 初始化颜色为灰色，用于背景点

        if len(points_for_clustering) >= 100 and self.enable_clustering:
            # 如果ROI内有点云，且启用了聚类
            pcd_temp = o3d.geometry.PointCloud() # 创建一个临时的Open3D点云对象
            pcd_temp.points = o3d.utility.Vector3dVector(points_for_clustering[:, :3])  # 设置点云的坐标
            downsampled_pcd = pcd_temp.voxel_down_sample(voxel_size=self.voxel_size)  # 体素降采样
            labels = np.array(downsampled_pcd.cluster_dbscan(eps=self.cluster_eps, min_points=self.cluster_min_points, print_progress=False))  # 使用DBSCAN进行聚类
            
            pcd_tree = o3d.geometry.KDTreeFlann(downsampled_pcd)  # 创建KDTree
            full_labels = np.ones(len(points_for_clustering), dtype=np.int32) * -1 # 初始化标签为-1 (表示噪声点)
            for i in range(len(points_for_clustering)): # 将聚类标签映射回原始点云
                [k, idx, _] = pcd_tree.search_knn_vector_3d(points_for_clustering[i, :3], 1) # k最近邻搜索
                if idx[0] < len(labels):
                    full_labels[i] = labels[idx[0]]
                
            max_label = full_labels.max() # 获取最大的聚类标签值
            if max_label > -1: self.get_logger().info(f"在ROI内检测到 {max_label + 1} 个聚类") # 记录聚类数量
            
            colors_roi = np.zeros((len(points_for_clustering), 3))  # 为ROI内的点分配颜色
            for label in range(-1, max_label + 1):
                label_mask = full_labels == label
                if label == -1:
                    colors_roi[label_mask] = [0.5, 0.5, 0.5] # 噪声点（灰色）
                else:
                    hue = label / (max_label + 1) if max_label > -1 else 0 # 计算色调值
                    colors_roi[label_mask] = self.hsv_to_rgb(hue, 0.8, 0.9)  # 使用HSV转RGB函数生成颜色
            
            # 将ROI内的颜色应用到最终的颜色数组中
            final_colors[roi_mask] = colors_roi # 设置ROI内点的颜色
            final_colors[~roi_mask] = [0.3, 0.3, 0.3] # 设置背景点的颜色为深灰色

        # 步骤4: 发布处理后的点云
        pcl_msg = create_pointcloud2_msg(msg.header, final_points) # 创建点云消息
        self.pcl_pub.publish(pcl_msg)  # 发布点云消息
        
        # 步骤5: 更新Open3D 可视化
        self.pcd.points = o3d.utility.Vector3dVector(final_points[:, :3]) # 设置Open3D点云的坐标
        self.pcd.colors = o3d.utility.Vector3dVector(final_colors)  # 设置Open3D点云的颜色
        
        if self.is_first_frame:
            # 如果是第一帧点云，需要添加点云到可视化窗口
            if len(self.pcd.points) > 0:  # 避免添加空点云
                self.vis.add_geometry(self.pcd)  # 将点云添加到可视化窗口
                self.get_logger().info(f"成功渲染第一帧，包含 {len(self.pcd.points)} 个点。")
                ctr = self.vis.get_view_control() # 获取视图控制器
                ctr.set_up([-10.0, 0.0, 10.0]) # Z轴为上
                ctr.set_front([-1, 0, -1]) # 相机从X轴负方向朝前看
                ctr.set_lookat([0, 0, -20]) # 相机看向原点
                ctr.set_zoom(0.2)
                self.is_first_frame = False  # 标记已添加过第一帧
        else:
            # 如果不是第一帧，则更新点云数据
            self.vis.update_geometry(self.pcd)
        
        # 轮询事件并更新渲染器，确保GUI能够响应
        self.vis.poll_events()  # 轮询事件，处理用户输入等
        self.vis.update_renderer()  # 更新渲染，显示新的点云数据

        # 统计处理时间，并记录日志
        if self.save_enabled:
            self.save_current_frame() # 保存点云（如果启用）
        self.frame_count += 1 # 增加帧计数
        process_time = (time.time() - start_time) * 1000 # 计算处理时间（毫秒）
        self.get_logger().info(f"帧 {self.frame_count}: 原始{num_original_points}点 → 过滤后{len(filtered_points)}点 → 处理时间{process_time:.2f}ms")

    def hsv_to_rgb(self, h, s, v):
        """
        HSV颜色空间转RGB颜色空间。
        用于生成鲜艳的聚类颜色。
        """
        if s == 0.0:
            return [v, v, v]  # 如果饱和度为0，返回灰度值
        
        i = int(h * 6.0)
        f = (h * 6.0) - i
        p = v * (1.0 - s)
        q = v * (1.0 - s * f)
        t = v * (1.0 - s * (1.0 - f))
        i %= 6
        
        if i == 0: return [v, t, p]
        if i == 1: return [q, v, p]
        if i == 2: return [p, v, t]
        if i == 3: return [p, q, v]
        if i == 4: return [t, p, v]
        if i == 5: return [v, p, q]

    def save_current_frame(self):
        """
        保存当前帧的点云数据到文件。
        """
        filename = f"{self.save_path}/pointcloud_{self.frame_count:06d}.pcd"  # 构建文件名
        o3d.io.write_point_cloud(filename, self.pcd)  # 保存点云
        self.get_logger().info(f"已保存点云: {filename}")  # 记录日志，表示已保存

    def destroy_node(self):
        """
        节点关闭时的清理工作。
        """
        self.get_logger().info("正在关闭节点...")
        self.vis.destroy_window()  # 销毁Open3D可视化窗口
        super().destroy_node()  # 调用父类的destroy_node方法
        self.get_logger().info("节点已关闭。")

def main(args=None):
    rclpy.init(args=args)  # 初始化rclpy
    node = LidarProcessor()  # 创建 LidarProcessor 节点
    try:
        # 节点运行, 持续处理ROS事件（如消息接收）
        rclpy.spin(node)
    except KeyboardInterrupt:
        # 捕获 Ctrl+C 信号，用于优雅地退出
        node.get_logger().info("捕获到Ctrl+C，开始关闭...")
    finally:
        # 确保节点和rclpy能够被安全关闭
        node.destroy_node()  # 销毁节点
        rclpy.shutdown()  # 关闭rclpy
if __name__ == '__main__':
    main()
